/**
 * Copyright (C) 2018-2021
 * All rights reserved, Designed By www.yixiang.co
 */
package co.yixiang.modules.system.utils;

import cn.hutool.json.JSONUtil;
import com.volcengine.ark.runtime.service.ArkService;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author: ZhangHouYing
 * @date: 2019-08-24 15:44
 */
@ServerEndpoint(value = "/websocket/ai")
@Component
public class WebSocket {

    private static ConcurrentHashMap<String, WebSocket> webSocketMap = new ConcurrentHashMap<>();
    //实例一个session，这个session是websocket的session
    private Session session;

    //新增一个方法用于主动向客户端发送消息
    public static void sendMessage(Object message, String userId) {
        WebSocket webSocket = webSocketMap.get(userId);
        if (webSocket != null) {
            try {
                webSocket.session.getBasicRemote().sendText(JSONUtil.toJsonStr(message));
//				System.out.println("【websocket消息】发送消息成功,用户="+userId+",消息内容"+message.toString());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public static ConcurrentHashMap<String, WebSocket> getWebSocketMap() {
        return webSocketMap;
    }

    public static void setWebSocketMap(ConcurrentHashMap<String, WebSocket> webSocketMap) {
        WebSocket.webSocketMap = webSocketMap;
    }

    //前端请求时一个websocket时
    @OnOpen
    public void onOpen(Session session) {
        this.session = session;
        webSocketMap.put(session.getId(), this);
        Map map = new HashMap();
        map.put("code", 200);
        map.put("msg", "连接成功");
        sendMessage(JSONUtil.toJsonStr(map), session.getId());
        System.out.println("【websocket消息】有新的连接,连接id" + session.getId());
    }

    //前端关闭时一个websocket时
    @OnClose
    public void onClose(Session session) {
        webSocketMap.remove(session.getId());
        System.out.println("【websocket消息】连接断开,总数:" + webSocketMap.size());
    }

    //前端向后端发送消息
    @OnMessage
    public void onMessage(Session session, String message) {
        if (!message.equals("ping")) {
            System.out.println("【websocket消息】收到客户端发来的消息:" + message);
        }
        ArkService service = null;
        try {
            service = ChatCompletionsExample.question(session.getId(), message);
        } catch (Exception e) {
            sendMessage("服务器连接超时，请稍后再试", session.getId());
            e.printStackTrace();
        } finally {
            if (service != null) {
                service.shutdownExecutor();
            }
        }
    }

}
